import cv2
import numpy as np
import math


def rgb2hsi(rgb: np.ndarray):
    """
    RGB颜色空间变HSI颜色空间
    :param rgb: 输入张量格式H*W*3，顺序RGB，数据类型numpy ndarray，范围0-255 uint8 或 0-1 float
    :return: 输出元组包含三个张量，格式 H*W，顺序HSI，数据类型0-1 float
    """
    if rgb.dtype == 'uint8':
        rgb = rgb.copy() / 255
    r, g, b = rgb[:, :, 0], rgb[:, :, 1], rgb[:, :, 2]
    theta = np.arccos(
        np.divide(0.5 * ((r - g) + (r - b)), ((r - g) ** 2 + (r - b) * (g - b)) ** 0.5 + 1e-6))
    h = np.where(b < g, theta, 2 * np.pi - theta) / (2 * np.pi)
    s = 1 - 3 * np.min(np.stack((r, g, b), axis=2), axis=2) / (r + g + b)
    i = (r + g + b) / 3
    h = np.clip(h, 0, 1)
    s = np.clip(s, 0, 1)
    i = np.clip(i, 0, 1)
    return h, s, i


def hsi2rgb(hsi: np.ndarray):
    """
    HSI颜色空间变RGB颜色空间
    输入张量格式H*W*3，顺序HSI，数据类型0-1 float
    输出张量格式H*W*3，顺序RGB，数据类型numpy ndarray，范围0-255 uint8
    """
    hsi_height, hsi_width, hsi_channel = hsi.shape
    r = np.zeros_like(hsi[:, :, 0])
    g = np.zeros_like(r)
    b = np.zeros_like(r)
    hsi = hsi.copy()
    h, s, i = hsi[:, :, 0], hsi[:, :, 1], hsi[:, :, 2]
    h = h * 2 * np.pi
    for m in range(0, hsi_height):
        for n in range(0, hsi_width):
            if 0 <= h[m, n] < 2 * np.pi / 3:
                b[m, n] = i[m, n] * (1 - s[m, n])
                r[m, n] = i[m, n] * (1 + s[m, n] * np.cos(h[m, n]) / np.cos(np.pi / 3 - h[m, n]))
                g[m, n] = 3 * i[m, n] - (r[m, n] + b[m, n])
            elif 2 * np.pi / 3 <= h[m, n] < 4 * np.pi / 3:
                h[m, n] = h[m, n] - 2 * np.pi / 3
                r[m, n] = i[m, n] * (1 - s[m, n])
                g[m, n] = i[m, n] * (1 + s[m, n] * np.cos(h[m, n]) / np.cos(np.pi / 3 - h[m, n]))
                b[m, n] = 3 * i[m, n] - (r[m, n] + g[m, n])
            else:
                h[m, n] = h[m, n] - 4 * np.pi / 3
                g[m, n] = i[m, n] * (1 - s[m, n])
                b[m, n] = i[m, n] * (1 + s[m, n] * np.cos(h[m, n]) / np.cos(np.pi / 3 - h[m, n]))
                r[m, n] = 3 * i[m, n] - (g[m, n] + b[m, n])
    r = np.clip(r, 0, 1)
    g = np.clip(g, 0, 1)
    b = np.clip(b, 0, 1)
    rgb = np.stack((r, g, b), axis=2)
    rgb = (rgb * 255).astype(np.uint8)
    return rgb


if __name__ == '__main__':
    src = cv2.imread('1.jpg')  # H*W*C
    print(src.shape)
    cv2.imshow("0", src)
    cv2.waitKey(0)

    src = cv2.cvtColor(src, cv2.COLOR_BGR2RGB)
    print(src.shape)
    # print(src[:10,:10,0])
    src = src / 255
    # print(src)
    src_height, src_width, src_channel = src.shape
    h, s, i = rgb2hsi(src)

    i = np.clip(i * 0.5, 0, 1)
    print(np.stack((h, s, i), axis=2).shape)
    rgb = hsi2rgb(np.stack((h, s, i), axis=2))

    # print(res[:10,:10,0])
    cv2.imshow("1", rgb[:, :, ::-1])  # H*W*C
    cv2.waitKey(0)
